"""
model name : MY_YOLO
file       : dataset.py
information:
    author : OuYang
    time   : 2025/1/22
"""
import os
import cv2
import torch
from torchvision import transforms
from torch.utils.data import Dataset


# YOLO Dataset
class YOLODataset(Dataset):
    def __init__(
            self,
            root,
            num_classes=3,
            imgsz=448,
            s=7,
            b=2,
            train=True,
            transform=None
    ):
        super(YOLODataset, self).__init__()
        self.root = root
        self.num_classes = num_classes
        self.imgsz = imgsz
        self.s = s
        self.b = b
        self.train = train
        self.transform = transform

        self.cell_size = 1.0 / self.s

        # Train Data or Valid Data
        if self.train:
            self.image_dir = os.path.join(self.root, 'train/images')
            self.label_dir = os.path.join(self.root, 'train/labels')
        else:
            self.image_dir = os.path.join(self.root, 'valid/images')
            self.label_dir = os.path.join(self.root, 'valid/labels')

        # Get Images
        self.images = os.listdir(self.image_dir)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        # from utils.show import show_image, show_image_from_tensor
        image_path = os.path.join(self.image_dir, self.images[index])
        label_path = os.path.join(self.label_dir, self.images[index].split('.')[0] + '.txt')

        # Open Image
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # print(f"Open Image: {image.shape}")
        # show_image(image)

        # Load Label
        labels = []
        with open(label_path, 'r') as file:
            for line in file:
                class_id, x, y, w, h = line.strip().split(' ')
                labels.append([
                    int(class_id),
                    float(x),
                    float(y),
                    float(w),
                    float(h)
                ])

        # Transform
        h, w = image.shape[:2]
        top, bottom, left, right = 0, 0, 0, 0
        if h > w:
            pad = (h - w) // 2
            left = pad
            right = pad
            for label in labels:
                label[1] = (label[1] * w + pad) / h
                label[3] = (label[3] * w) / h
        elif w > h:
            pad = (w - h) // 2
            top = pad
            bottom = pad
            for label in labels:
                label[2] = (label[2] * h + pad) / w
                label[4] = (label[4] * h) / w
        padded_image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])
        # print(f"Padded Image: {padded_image.shape}")
        # show_image(padded_image)
        resized_image = cv2.resize(padded_image, (self.imgsz, self.imgsz))
        # print(f"Resized image: {resized_image.shape}")
        # show_image(resized_image)
        # image = np.transpose(resized_image, (2, 0, 1))

        if self.transform is None:
            self.transform = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        image = self.transform(resized_image)
        # print(f"Transformed image: {image.shape}")
        # show_image_from_tensor(image)

        # Create Target
        target = self._build_target(labels)

        return image, target

    def _build_target(self, labels):
        target = torch.zeros((self.s, self.s, self._cal_one_grid_predictions()), dtype=torch.float32)

        for label in labels:
            class_id, x, y, w, h = label
            i, j = int(x // self.cell_size), int(y // self.cell_size)
            x = x / self.cell_size - i
            y = y / self.cell_size - j

            for bi in range(self.b):
                target[i][j][bi * 5: bi * 5 + 4] = torch.tensor([x, y, w, h], dtype=torch.float32)
                target[i][j][bi * 5 + 4] = 1
            target[i][j][self.b * 5 + class_id] = 1

        return target

    def _cal_one_grid_predictions(self):
        return self.b * 5 + self.num_classes


if __name__ == '__main__':
    train_set = YOLODataset(
        root='data/test',
        num_classes=20
    )
    valid_set = YOLODataset(
        root='data/test',
        num_classes=20,
        train=False
    )
    for i in range(5):
        image, target = train_set[i]

        # import numpy as np
        # print(image.shape)

        # image = np.array(image)
        # image = np.transpose(image, (1, 2, 0))
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        from utils.show import show_image_from_tensor, draw_boxes_from_tensor
        from utils.utils import xywh2xyxy

        # show_image_from_tensor(image)

        cell_size = 1.0 / 7
        boxs = []
        for i in range(7):
            for j in range(7):
                if target[i][j][4] == 1:
                    print(i, j)
                    x, y, w, h = target[i][j][:4]
                    x = x * cell_size + i * cell_size
                    y = y * cell_size + j * cell_size

                    x *= 448
                    y *= 448
                    w *= 448
                    h *= 448

                    x1, y1, x2, y2 = xywh2xyxy((x.item(), y.item(), w.item(), h.item()))
                    boxs.append([x1, y1, x2, y2])
                    print(x, y, w, h)

        image = draw_boxes_from_tensor(image, boxs, show=True)

